#include "pch.h"
#include "ScClient.h"

#define DATA_BUFSIZE 4096

#pragma pack(1)
typedef struct {
	TCHAR domain[128];
	TCHAR port[16];
} ScServerConfig;
#pragma pack()

ScServerConfig gScServerConfig;
ScClient *gScClientCoor = NULL;
ScClient *gScClientGame = NULL;

ScClient::ScClient() {	
	this->connSocket = INVALID_SOCKET;
	this->addrInfo = NULL;
	// Initialize the critical section one time only.
	if (!InitializeCriticalSectionAndSpinCount(&this->criticalSection, 0x00000400))
		MK::Logger::Instance()->Trace(TEXT("InitializeCriticalSectionAndSpinCount failed."));
}

ScClient::~ScClient() {
	// Release resources used by the critical section object.
	DeleteCriticalSection(&this->criticalSection);
}

void ScClient::release() {
	WSACloseEvent(this->sendOverlapped.hEvent);
	for (int i = 0; i < sizeof(this->recvOverlapped) / sizeof(WSAOverlappedWrapper); i++) {
		WSAOverlappedWrapper *wp = &this->recvOverlapped[i];
		WSACloseEvent(wp->overlapped.hEvent);
	}
	
	if (this->connSocket != INVALID_SOCKET) {
		closesocket(this->connSocket);
	}
	if (this->addrInfo != NULL) {
		freeaddrinfo(this->addrInfo);
	}

	WSACleanup();
}

BOOL ScClient::startConnection(const char *serverName, const char *serverPort) {
	WSADATA wsd;
	struct addrinfo *ptr = NULL, hints;
	int rc, err = 0;

	// Load Winsock
	rc = WSAStartup(MAKEWORD(2, 2), &wsd);
	if (rc != 0) {
		MK::Logger::Instance()->Trace(TEXT("Unable to load Winsock: %d"), rc);
		return FALSE;
	}
	// Make sure the hints struct is zeroed out
	SecureZeroMemory((PVOID)&hints, sizeof(struct addrinfo));

	// Initialize the hints to retrieve the server address for IPv4
	hints.ai_family = AF_INET;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;

	if (serverName == NULL) {
		serverName = this->serverName;
		serverPort = this->serverPort;
	}
	rc = getaddrinfo(serverName, serverPort, &hints, &this->addrInfo);
	if (rc != 0) {
		MK::Logger::Instance()->Trace(TEXT("getaddrinfo failed with error: %d"), rc);
		return FALSE;
	}

	for (ptr = this->addrInfo; ptr != NULL; ptr = ptr->ai_next) {
		this->connSocket = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
		if (this->connSocket == INVALID_SOCKET) {
			MK::Logger::Instance()->Trace(TEXT("socket failed with error: %d"), WSAGetLastError());
			freeaddrinfo(this->addrInfo);
			this->addrInfo = NULL;
			return FALSE;
		}

		rc = connect(this->connSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
		if (rc == SOCKET_ERROR) {
			if (WSAECONNREFUSED == (err = WSAGetLastError())) {
				closesocket(this->connSocket);
				this->connSocket = INVALID_SOCKET;
				continue;
			}
			MK::Logger::Instance()->Trace(TEXT("connect failed with error: %d"), err);
			freeaddrinfo(this->addrInfo);
			closesocket(this->connSocket);
			this->addrInfo = NULL;
			this->connSocket = INVALID_SOCKET;
			return FALSE;
		}
		break;
	}

	if (this->connSocket == INVALID_SOCKET) {
		MK::Logger::Instance()->Trace(TEXT("Unable to establish connection with the server!"));
		freeaddrinfo(this->addrInfo);
		this->addrInfo = NULL;
		return FALSE;
	}

	// Create an event handle and setup the overlapped structure.
	SecureZeroMemory((PVOID)&this->sendOverlapped, sizeof(WSAOVERLAPPED));
	this->sendOverlapped.hEvent = WSACreateEvent();
	if (this->sendOverlapped.hEvent == NULL) {
		MK::Logger::Instance()->Trace(TEXT("WSACreateEvent failed with error: %d"), WSAGetLastError());
		freeaddrinfo(this->addrInfo);
		closesocket(this->connSocket);
		this->addrInfo = NULL;
		this->connSocket = INVALID_SOCKET;
		return FALSE;
	}

	// Create an event handle and setup an overlapped structure.
	SecureZeroMemory((PVOID)&this->recvOverlapped, sizeof(this->recvOverlapped));
	for (int i = 0; i < sizeof(this->recvOverlapped) / sizeof(WSAOverlappedWrapper); i++) {
		WSAOverlappedWrapper *wp = &this->recvOverlapped[i];
		wp->overlapped.hEvent = WSACreateEvent();
		if (wp->overlapped.hEvent == NULL) {
			MK::Logger::Instance()->Trace(TEXT("WSACreateEvent failed with error: %d"), WSAGetLastError());
			freeaddrinfo(this->addrInfo);
			closesocket(this->connSocket);
			this->addrInfo = NULL;
			this->connSocket = INVALID_SOCKET;
			return FALSE;
		}
	}

	TCHAR wServerName[64], wServerPort[8];
	mbstowcs_s(NULL, wServerName, serverName, sizeof(wServerName) / sizeof(TCHAR));
	mbstowcs_s(NULL, wServerPort, serverPort, sizeof(wServerPort) / sizeof(TCHAR));
	MK::Logger::Instance()->Trace(TEXT("Client connected to (%s:%s)"), wServerName, wServerPort);

	if (this->serverName != serverName) {
		strcpy_s(this->serverName, sizeof(this->serverName), serverName);
	}
	if (this->serverPort != serverPort) {
		strcpy_s(this->serverPort, sizeof(this->serverPort), serverPort);
	}

	return TRUE;
}

BOOL ScClient::assignDataBuf(const char *buf, int len, WSABUF **pDataBuf, int &bufCnt, DWORD targetIp) {
	WSABUF *dataBuf, *dBuf;
	bufCnt = (len + sizeof(DWORD) + DATA_BUFSIZE - 1) / DATA_BUFSIZE;
	dataBuf = new WSABUF[bufCnt];
	if (dataBuf == NULL) {
		return FALSE;
	}
	memset(dataBuf, 0x0, sizeof(WSABUF) * bufCnt);
	*pDataBuf = dataBuf;
	
	for (int i = 0; i < bufCnt - 1; i++) {
		dBuf = &dataBuf[i];
		dBuf->buf = new char[DATA_BUFSIZE];
		if (dBuf->buf == NULL) {
			return FALSE;
		}
		memcpy(dBuf->buf, &buf[i * DATA_BUFSIZE - sizeof(DWORD)], DATA_BUFSIZE);
		dBuf->len = DATA_BUFSIZE;
	}
	if (len + (int)sizeof(DWORD) > (bufCnt - 1) * DATA_BUFSIZE) {
		dBuf = &dataBuf[bufCnt - 1];
		dBuf->buf = new char[DATA_BUFSIZE];
		if (dBuf->buf == NULL) {
			return FALSE;
		}
		dBuf->len = len + sizeof(DWORD) - (bufCnt - 1) * DATA_BUFSIZE;
		memcpy(dBuf->buf, &buf[(bufCnt - 1) * DATA_BUFSIZE - sizeof(DWORD)], dBuf->len);
	}

	DWORD *ip = (DWORD *)(dataBuf[0].buf);
	*ip = targetIp;

	return TRUE;
}

void ScClient::releaseDataBuf(WSABUF **pDataBuf, int bufCnt) {
	WSABUF *dataBuf = *pDataBuf, *dBuf;

	for (int i = 0; i < bufCnt; i++) {
		dBuf = &dataBuf[i];
		if (dBuf->buf != NULL) {
			delete dBuf->buf;
		}
	}
	delete dataBuf;
	*pDataBuf = NULL;
}

int ScClient::send(const char *buf, int len, DWORD targetIp) {
	WSABUF dataBuf, *pDataBuf;
	DWORD sendBytes, flags;
	char buffer[DATA_BUFSIZE];
	int err = 0, curDataBufCnt;

	if (len <= DATA_BUFSIZE - sizeof(UINT32)) {
		dataBuf.len = len + sizeof(DWORD);
		dataBuf.buf = buffer;
		DWORD *ip = (DWORD *)buffer;
		*ip = targetIp;
		memcpy(&buffer[sizeof(DWORD)], buf, len);

		pDataBuf = &dataBuf;
		curDataBufCnt = 1;
	} else {
		if (!this->assignDataBuf(buf, len, &pDataBuf, curDataBufCnt, targetIp)) {
			this->releaseDataBuf(&pDataBuf, curDataBufCnt);
			return -1;
		}
	}
	
	int rc, retry = 0, maxRetry = 3;
	BOOL bReconnect = FALSE;
	while (true)
	{
		if (bReconnect) {
			bReconnect = FALSE;
			MK::Logger::Instance()->Trace(TEXT("Try to reconnect, retry(%d) ..."), retry);
			this->release();
			this->startConnection(NULL, NULL);
		}
		rc = WSASend(this->connSocket, pDataBuf, curDataBufCnt, &sendBytes, 0, &this->sendOverlapped, NULL);
		if ((rc == SOCKET_ERROR) && (WSA_IO_PENDING != (err = WSAGetLastError()))) {
			MK::Logger::Instance()->Trace(TEXT("WSASend failed with error: %d"), err);
			if (retry++ < maxRetry) {
				bReconnect = TRUE;
				continue;
			}
			else {
				return -1;
			}
		}

		rc = WSAWaitForMultipleEvents(1, &this->sendOverlapped.hEvent, TRUE, INFINITE, TRUE);
		if (rc == WSA_WAIT_FAILED) {
			MK::Logger::Instance()->Trace(TEXT("WSAWaitForMultipleEvents failed with error: %d"), WSAGetLastError());
			if (retry++ < maxRetry) {
				bReconnect = TRUE;
				continue;
			}
			else {
				return -1;
			}
		}

		rc = WSAGetOverlappedResult(this->connSocket, &this->sendOverlapped, &sendBytes, FALSE, &flags);
		if (rc == FALSE) {
			MK::Logger::Instance()->Trace(TEXT("WSASend failed with error: %d"), WSAGetLastError());
			if (retry++ < maxRetry) {
				bReconnect = TRUE;
				continue;
			}
			else {
				return -1;
			}
		}

		// MK::Logger::Instance()->Trace(TEXT("ScClient send %d bytes"), sendBytes);
		WSAResetEvent(this->sendOverlapped.hEvent);
		break;
	}

	if (curDataBufCnt > 1) {
		this->releaseDataBuf(&pDataBuf, curDataBufCnt);
	}

	return sendBytes;
}

WSAOVERLAPPED *ScClient::getRecvOverLapped(DWORD threadId) {
	WSAOVERLAPPED *overlapped = NULL;
	EnterCriticalSection(&this->criticalSection);
	
	for (int i = 0; i < sizeof(this->recvOverlapped) / sizeof(WSAOverlappedWrapper); i++) {
		WSAOverlappedWrapper *wp = &this->recvOverlapped[i];
		if (wp->threadId == threadId) {
			overlapped = &wp->overlapped;
			break;
		}
		else if (wp->threadId == 0) {
			wp->threadId = threadId;
			overlapped = &wp->overlapped;
			break;
		}
	}
	
	LeaveCriticalSection(&this->criticalSection);

	return overlapped;
}

void ScClient::clearRecvOverLappedThreadId() {
	WSAOVERLAPPED *overlapped = NULL;
	EnterCriticalSection(&this->criticalSection);

	for (int i = 0; i < sizeof(this->recvOverlapped) / sizeof(WSAOverlappedWrapper); i++) {
		WSAOverlappedWrapper *wp = &this->recvOverlapped[i];
		wp->threadId = 0;
	}

	LeaveCriticalSection(&this->criticalSection);
}

int ScClient::recv(LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, DWORD threadId, DWORD &fromIp) {
	DWORD flags = 0;
	int rc, err = 0, retry = 0, maxRetry = 3;
	BOOL bReconnect = FALSE;

	WSAOVERLAPPED *overlapped = this->getRecvOverLapped(threadId);
	if (overlapped == NULL) {
		MK::Logger::Instance()->Trace(TEXT("Cannot get recv overlapped"));
		return SOCKET_ERROR;
	}

	while (true)
	{
		if (bReconnect) {
			bReconnect = FALSE;
			MK::Logger::Instance()->Trace(TEXT("Try to reconnect, retry(%d) ..."), retry);
			this->release();
			this->startConnection(NULL, NULL);
		}
		rc = WSARecv(this->connSocket, lpBuffers, dwBufferCount, lpNumberOfBytesRecvd, &flags, overlapped, NULL);
		if ((rc == SOCKET_ERROR) && (WSA_IO_PENDING != (err = WSAGetLastError()))) {
			MK::Logger::Instance()->Trace(TEXT("WSARecv failed with error: %d"), err);
			return SOCKET_ERROR;
		}

		rc = WSAWaitForMultipleEvents(1, &overlapped->hEvent, TRUE, INFINITE, TRUE);
		if (rc == WSA_WAIT_FAILED) {
			MK::Logger::Instance()->Trace(TEXT("WSAWaitForMultipleEvents failed with error: %d"), WSAGetLastError());
			return SOCKET_ERROR;
		}

		rc = WSAGetOverlappedResult(this->connSocket, overlapped, lpNumberOfBytesRecvd, FALSE, &flags);
		if (rc == FALSE) {
			MK::Logger::Instance()->Trace(TEXT("WSARecv operation failed with error: %d"), WSAGetLastError());
			return SOCKET_ERROR;
		}

		// MK::Logger::Instance()->Trace(TEXT("Read %d bytes"), *lpNumberOfBytesRecvd);
		WSAResetEvent(overlapped->hEvent);

		if (*lpNumberOfBytesRecvd > 0) {
			break;
		}
		if (retry++ < maxRetry) {
			continue;
		} else {
			bReconnect = TRUE;
			retry = 0;
		}
	}

	if (*lpNumberOfBytesRecvd > 0) {
		fromIp = *((DWORD *)lpBuffers->buf);
		*lpNumberOfBytesRecvd = *lpNumberOfBytesRecvd - sizeof(DWORD);
		int bufLen, appendLen, restLen = *lpNumberOfBytesRecvd;
		for (int i = 0; i < (int)dwBufferCount; i++) {
			char *buf = lpBuffers[i].buf;
			bufLen = (restLen > (int)lpBuffers[i].len - (int)sizeof(DWORD))? ((int)lpBuffers[i].len - (int)sizeof(DWORD)) : restLen;
			memcpy(buf, &lpBuffers->buf[sizeof(DWORD)], bufLen);
			restLen -= bufLen;
			if (i < (int)dwBufferCount - 1) {
				appendLen = (restLen > (int)sizeof(DWORD)) ? (int)sizeof(DWORD) : restLen;
				memcpy(&buf[bufLen], lpBuffers[i + 1].buf, appendLen);
				restLen -= appendLen;
			}
		}
		
	}
	
	return 0;
}

void DLL_EXPORT setScServerConfig(LPVOID params) {
	memcpy(&gScServerConfig, params, sizeof(ScServerConfig));
	MK::Logger::Instance()->Trace(TEXT("setScServerConfig (%s:%s)"), gScServerConfig.domain, gScServerConfig.port);

	TCHAR *delimitor = wcschr(gScServerConfig.port, L',');
	delimitor[0] = L'\0';

	char domain[128], port_coor[8], port_game[8];
	wcstombs_s(NULL, domain, gScServerConfig.domain, sizeof(domain));
	wcstombs_s(NULL, port_coor, gScServerConfig.port, sizeof(port_coor));
	wcstombs_s(NULL, port_game, &delimitor[1], sizeof(port_game));
	
	gScClientCoor = new ScClient();
	if (!gScClientCoor->startConnection(domain, port_coor)) {
		gScClientCoor->release();
		delete gScClientCoor;
		gScClientCoor = NULL;
	}

	gScClientGame = new ScClient();
	if (!gScClientGame->startConnection(domain, port_game)) {
		gScClientGame->release();
		delete gScClientGame;
		gScClientGame = NULL;
	}
}